from python_ai.CV_4.single_respect.darknet19.darknet19_pytorch_master.my_try_pretrained_darknet19 import MyDarknet19
import torch as pt
import numpy as np
import matplotlib.pyplot as plt
from python_ai.common.xcommon import *


class RearNet(pt.nn.Module):

    def __init__(self, in_ch, n_cls, **kwargs):
        super().__init__(**kwargs)
        self.avgpool = pt.nn.AdaptiveAvgPool2d((1, 1))
        self.fc = pt.nn.Linear(in_ch, n_cls)

    def forward(self, x):
        x = self.avgpool(x)
        x = pt.squeeze(x, dim=3)
        x = pt.squeeze(x, dim=2)
        x = self.fc(x)
        return x


if '__main__' == __name__:
    from python_ai.common.read_data.cat_dog_teachers import load_data
    from python_ai.common.routine.pt_get_device import get_device
    from python_ai.common.routine.pt_compile_fit_evaluate import *

    BATCH_SIZE = 16
    N_EPOCHS = 128
    ALPHA = 1e-5
    VER = 'v1.0'
    WEIGHT_PATH = '../../../../../../large_data/model/darknet19/darknet19-deepBakSu-e1b3ec1e.pth'
    data_path = r'../../../../../../large_data/DL1/_many_files/zoo'
    IMG_H = 256
    IMG_W = 256

    FILE_NAME = os.path.basename(__file__)
    LOG_DIR = os.path.join('_log', FILE_NAME, VER)

    device = get_device()

    (x_train, y_train), (x_val, y_val), (x_test, y_test) = load_data(data_path, 1, IMG_H, IMG_W)

    net = MyDarknet19(WEIGHT_PATH, in_ch=3, n_cls=1000, is_pretrained=True).features.to(device)
    out_channels = net[22].conv.out_channels

    x_train = trans_data(x_train, net, device, BATCH_SIZE)
    x_test = trans_data(x_test, net, device, BATCH_SIZE)
    x_val = trans_data(x_val, net, device, BATCH_SIZE)

    print('x_train', x_train.shape)
    print('x_val', x_val.shape)
    print('x_test', x_test.shape)

    x_train = pt.Tensor(x_train)
    x_test = pt.Tensor(x_test)
    x_val = pt.Tensor(x_val)
    y_train = pt.Tensor(y_train)
    y_test = pt.Tensor(y_test)
    y_val = pt.Tensor(y_val)

    ds_train = TensorDataset(x_train, y_train)
    ds_test = TensorDataset(x_test, y_test)
    ds_val = TensorDataset(x_val, y_val)

    dl_train = DataLoader(ds_train, BATCH_SIZE, drop_last=True)
    dl_test = DataLoader(ds_test, BATCH_SIZE, drop_last=False)
    dl_val = DataLoader(ds_val, BATCH_SIZE, drop_last=False)

    model = RearNet(out_channels, 2).to(device)
    criterion = pt.nn.CrossEntropyLoss()

    model_dict = compile(device, model, criterion, pt.optim.Adam, ALPHA, accuracy)
    loss_his, acc_his, loss_his_val, acc_his_val = fit(model_dict, dl_train, dl_val, N_EPOCHS)

    evaluate(model_dict, dl_test)

    plt.figure(figsize=[12, 6])
    plt_dict = dict(spr=1, spc=2, spn=0)

    plot_loss_and_metric(plt_dict, loss_his, loss_his_val, acc_his, acc_his_val)

    plt.show()
